package com.breeze.components.socket.memory;

import com.breeze.components.socket.WebSocket;
import com.breeze.components.socket.WebSocketEvent;
import com.breeze.components.socket.WebSocketManager;
import com.breeze.components.socket.util.WebSocketUtil;
import com.breeze.components.socket.vo.WebSocketRequest;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;

/**
 * 内存websocket管理者
 *
 * @author breeze
 */
@Slf4j
@Getter
public class MemWebSocketManager implements WebSocketManager, ApplicationContextAware {

  private final Map<String, WebSocket> connections = new ConcurrentHashMap<>(100);
  private ApplicationContext applicationContext;

  @Override
  public void setApplicationContext(@NonNull ApplicationContext applicationContext)
      throws BeansException {
    this.applicationContext = applicationContext;
  }

  @Override
  public WebSocket get(String identifier) {
    return connections.get(identifier);
  }

  @Override
  public void put(String identifier, WebSocket webSocket) {
    connections.put(identifier, webSocket);
    // 发送连接事件
    getApplicationContext()
        .publishEvent(new WebSocketEvent(webSocket, WebSocketEvent.EVENT_TYPE_OPEN));
  }

  @Override
  public void remove(String identifier) {
    WebSocket removedWebSocket = connections.remove(identifier);
    // 发送关闭事件
    if (Objects.nonNull(removedWebSocket)) {
      removedWebSocket.closeSession();
      getApplicationContext()
          .publishEvent(new WebSocketEvent(removedWebSocket, WebSocketEvent.EVENT_TYPE_CLOSE));
    }
  }

  @Override
  public Map<String, WebSocket> localWebSocketMap() {
    return connections;
  }

  @Override
  public void sendMessage(String identifier, String message) {
    WebSocket webSocket = get(identifier);
    if (Objects.isNull(webSocket)) {
      return;
    }
    WebSocketUtil.sendMessage(webSocket.getSession(), message);
  }

  @Override
  public void broadcast(String message) {
    localWebSocketMap()
        .values()
        .forEach(webSocket -> WebSocketUtil.sendMessage(webSocket.getSession(), message));
  }

  @Override
  public void onMessage(String identifier, WebSocketRequest request) {
    WebSocket webSocket = connections.get(identifier);
    // 发送消息事件
    if (Objects.nonNull(webSocket)) {
      getApplicationContext()
          .publishEvent(new WebSocketEvent(webSocket, WebSocketEvent.EVENT_TYPE_MESSAGE, request));
    }
  }
}
